package union_find

// UF第二版：是一个树形结构，会把左边子树接上右边子树
type UnionFindV2 struct {
	parent []int
}

func (u *UnionFindV2) Size() int {
	return len(u.parent)
}

// 查找过程，查找元素p所对应的集合编号
// O(h)复杂度，为树的高度
func (u *UnionFindV2) find(p int) int {
	if p < 0 || p > len(u.parent) {
		panic("p is out of range")
	}
	// 不断的查询自己的父节点，直到到达根节点
	// 根节点的特点是 parent[p] == p

	for p != u.parent[p] {
		p = u.parent[p]
	}
	return p
}

func (u *UnionFindV2) IsConnected(p int, q int) bool {
	return u.find(p) == u.find(q)
}

func (u *UnionFindV2) Union(p int, q int) {
	pRoot := u.find(p)
	qRoot := u.find(q)

	if pRoot == qRoot {
		return
	}

	u.parent[pRoot] = qRoot
}

func NewUFV2(size int) *UnionFindV2 {
	parent := make([]int, size)
	for i := 0; i < len(parent); i++ {
		parent[i] = i
	}
	return &UnionFindV2{parent}
}
